# -*- coding:utf-8 -*-
# Author:wudb
# @File : main.py
# @Time : 20-6-15 下午3:54

import os

os.environ['TF_CPP_MIN_LOG_LEVEL'] = '2'
import tensorflow as tf

gpus = tf.config.experimental.list_physical_devices(device_type='GPU')
for gpu in gpus:
    tf.config.experimental.set_memory_growth(gpu, True)

os.environ['CUDA_VISIBLE_DEVICES'] = '2,3'

from tensorflow import keras
from tensorflow.keras import Sequential, metrics, layers, optimizers, datasets

def process(x, y):
    x = 2 * tf.cast(x, dtype= tf.float32)/255. - 1
    x = tf.reshape(x,[28*28])
    y = tf.cast(y, dtype = tf.int32)
    y = tf.one_hot(y, depth = 10 )
    return x,y
(x_train, y_train),(x_test, y_test) = datasets.mnist.load_data()

db_train = tf.data.Dataset.from_tensor_slices((x_train, y_train))
db_train = db_train.shuffle(1000).map(process).batch(256)

db_test = tf.data.Dataset.from_tensor_slices((x_test, y_test))
db_test = db_test.map(process).batch(256)

# [b, 28, 28] [b]
sample = next(iter(db_train))
print(sample[0].shape,tf.reduce_max(sample[0]).numpy(),tf.reduce_min(sample[0]).numpy(),
      sample[1].shape, tf.reduce_max(sample[1]).numpy(), tf.reduce_min(sample[1]).numpy())

model = Sequential([
    layers.Dense(512,activation='relu'),
    layers.Dense(256,activation='relu'),
    layers.Dense(128,activation='relu'),
    layers.Dense(64, activation='relu'),
    layers.Dense(10)
])

def main():
    model.build(input_shape= (None,28*28))
    model.summary()
    model.compile(optimizer = optimizers.Adam(lr=0.001),
              loss = tf.losses.CategoricalCrossentropy(from_logits=True),
              metrics = ['accuracy'])
    model.fit(db_train, epochs = 30, validation_data= db_test, validation_freq = 1)
    model.evaluate(db_test)

if __name__ == '__main__':
    main()